# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import json
import logging
from typing import List, Union

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.kraken.kr_rest.product import KrakenProduct


class KrakenSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)

    self._register_url('wss://ws.kraken.com')
    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('kraken_ws_depth_%s%s' % (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('kraken_ws_trade_%s%s' % (group_id, worker_id))[0]

    self._publisher = {
        'book': queue.get_writer(topic_id=topic_depth[1]),
        'trade': queue.get_writer(topic_id=topic_trade[1])
    }

    self._enable_topic_event_publisher(queue,
                                       'kraken_rest_topic_event', [topic_depth[0], topic_trade[0]])
    self._add_teardown_callback(self._teardown_publisher)

  @gen.coroutine
  def on_open(self):
    yield self._send_subscribe_request()

  @gen.coroutine
  def _send_subscribe_request(self):
    for symbol in self._symbol:
      product = KrakenProduct.FromStrNativeProduct(symbol)
      subscription_symbol = product.native_subscription_symbol
      self._logger.debug('Subscribing %s' % subscription_symbol)

      subscribe_depth_str = json.dumps({
          "event": "subscribe",
          "pair": [subscription_symbol],
          "subscription": {
              "name": "book", "depth": 25
          }
      })
      yield self._ws.write_message(subscribe_depth_str)

      subscribe_trade_str = json.dumps({
          "event": "subscribe", "pair": [subscription_symbol], "subscription": {
              "name": "trade",
          }
      })
      yield self._ws.write_message(subscribe_trade_str)
      yield gen.sleep(0.5)  # Sleep 0.5 second before next request

    self._publish_topic_event('START')
    self._logger.info('Subscribe done')

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: bytes):
    if self.stopped:
      return

    if isinstance(message, dict):
      event = message['event']
      if event == 'systemStatus':
        self._publish_topic_event(message)
      elif event == 'subscriptionStatus':
        event_type = message['subscription']['name']
        try:
          self._publisher[event_type].write(timestamp, message)
        except (KeyError, ValueError):
          self._logger.error('Unknown channel: %s' % message)
      else:
        # Probably heartbeat, ping pong, don't care.
        pass
    elif isinstance(message, list):
      # Book message is a list and second element is a dict.
      # [
      #   1234,
      #   {"b": [
      #     [
      #       "5541.30000",
      #       "0.00000000",
      #       "1534614335.345903"
      #     ]
      #   ]}
      # ]

      # Trade message is a list and second element is also a list.
      # [
      #   0,
      #   [
      #     [
      #       "5541.20000",
      #       "0.15850568",
      #       "1534614057.321597",
      #       "s",
      #       "l",
      #       ""
      #     ]
      #   ]
      # ]
      if isinstance(message[1], dict):
        self._publisher['book'].write(timestamp, message)
      elif isinstance(message[1], list):
        self._publisher['trade'].write(timestamp, message)
      else:
        self._logger.error('Unknown message: %s' % message)
    else:
      self._logger.error('Unknown message: %s' % message)
